#!/usr/bin/env python

import os

#
# Generators that split a file, directory or list-of-files into piece-sized pieces.
# Yields the data one piece at a time.
#

def create_pieces_file(piece_size, partial_piece, file):
    fd = open(file, "rb")
    n_bytes_to_read = piece_size - len(partial_piece)

    more_data=True
    while more_data:
        bytes = fd.read(n_bytes_to_read)
        n_bytes_to_read = piece_size
        if len(bytes) == 0:
            more_data = False
        else:
	    partial_piece += bytes
	if len(partial_piece) >= piece_size:
	    yield partial_piece
	    partial_piece = ""
    fd.close()
    if len(partial_piece) > 0:
        yield partial_piece
        partial_piece = ""

def create_pieces_files(piece_size, partial_piece, files):
    for name in files:
        for piece in create_pieces_file(piece_size, partial_piece, name):
	    if len(piece) >= piece_size:
                yield piece
                partial_piece = ""
            else:
                partial_piece = piece
    if len(partial_piece) > 0:
        yield partial_piece
        partial_piece = ""

def create_pieces_dir(piece_size, partial_piece, directory):
    all_files = []
    for root, dirs, files in os.walk(directory):
        all_files.extend( [ os.path.join(root, f) for f in files ] )
    for piece in create_pieces_files(piece_size, partial_piece, all_files):
        yield piece

if __name__ == '__main__':
    import sys
#    import pdb; pdb.set_trace()

    i=0
    pieceLength = 60
    if (os.path.isdir(sys.argv[1])):
        genfunc = create_pieces_dir
    else:
        genfunc = create_pieces_file

    for piece in genfunc(pieceLength, "", sys.argv[1]):
        print("Piece %d: %s"%(i, piece.replace("\n", " ")))
        if len(piece) > pieceLength:
	    print("ERROR! pieceLength is %d (> %d)"%(len(piece), pieceLength))
	i = i+1
